--[[
Module providing a simple class for Finite State Machines, which is a natural way to define things 
such as game menu navigation, simple AI for enemies, and many other applications.
--]]
module(..., package.seeall)


FSM = class("FSM")
FSM.enter_prefix      = "enter_"
FSM.exit_prefix       = "exit_"
FSM.transition_prefix = "transition_"
FSM.initial_state     = "init"

function FSM.__init(self)
    self.state = nil
    self.transitions = {ongoing=0, completed=0}
    self.driver = FSMDriver(self)
end

function FSM.transitioning(self)
    --[[ Returns true if the state machine is undergoing a state transition. ]]
    return self.transitions.ongoing > 0
end

function FSM.init(self, ...)
    --[[ Enter the machine's initial state. ]]
    return self:enter(self.initial_state, ...)
end

function FSM.enter(self, new_state, ...)
    --[[ Enter a given state in the machine. Before calling the enter function of the new state,
    the exit function of the old state is executed. Note that the exit function is optional. If 
    no exit function is found, the next state's enter function, if defined, will be called
    immediately. ]]
    if self.state ~= nil then
        self:exit()
    end
    local result
    if new_state ~= nil then
        self.state = new_state
        local enter_fnc = self:__get_function(self.enter_prefix, new_state)
        local transitions = self.transitions
        if enter_fnc ~= nil then
            transitions.ongoing = transitions.ongoing + 1
            result = enter_fnc(self, ...)
            transitions.ongoing = transitions.ongoing - 1
        end
        transitions.completed = transitions.completed + 1
    end
    return result
end

function FSM.exit(self)
    --[[ Exit the current state. This calls the current state's exit function (if one is defined) 
    and sets the FSM's state to nil if no new transition was made during execution of the exit 
    function. ]]
    local result
    local transitions = self.transitions
    local transition_no = transitions.completed
    local old_state = self.state
    if old_state ~= nil then
        local exit_fnc = self:__get_function(self.exit_prefix, old_state)
        if exit_fnc ~= nil then
            transitions.ongoing = transitions.ongoing + 1
            result = exit_fnc(self)
            transitions.ongoing = transitions.ongoing - 1
        end
        if transitions.completed == transition_no then
            self.state = nil
        end
    end
    return result
end

function FSM.input(self, symbol, ...)
    --[[ Feed a symbol to the state machine. The symbol is processed by the transition function 
    defined by the current state (or the default transition function if the current state does
    not have a transition function). The transition function should return the name of the 
    state the machine should jump to. ]]
    assert(self.state ~= nil, "machine state is undefined")
    local transition_fnc = self:__get_function(self.transition_prefix, self.state)
    assert(transition_fnc ~= nil, "unable to find transition function or map")
    local new_state
    if type(transition_fnc) == "table" then
        new_state = transition_fnc[symbol]
    else
        new_state = transition_fnc(self, symbol)
    end
    return self:enter(new_state, ...)
end

function FSM.__get_function(self, prefix, state)
    return self[prefix .. state] or self[prefix .. "default"]
end


--[[
The FSMDriver table allows feeding input symbols to the finite state machine as if they were 
methods of the driver object. For example:
    fsm.driver:run(60)
        is equivalent to
    fsm:input("run", 60)
Actually, attribute access on the driver object returns a partial method call to input() method on 
the FSM, where the input symbol is the attribute name. The remaining arguments are passed later 
when the partial object is actually called.
--]]
FSMDriver = class("FSMDriver")

function FSMDriver.__new(cls, fsm)
    return {fsm=fsm}
end

function FSMDriver.__index(self, key)
    return func.partial(self.fsm.input, self.fsm, key)
end

function FSMDriver.__newindex(self, key, value)
    error("illegal operation on fsm driver (read-only object)", 2)
end


--------------------------------------------------------------------------------
-- Test code
--------------------------------------------------------------------------------
function test()
    local foo = class("FooSM", FSM)
    foo.initial_state = "a"

    function foo.enter_default(self)
        print("entering state", self.state)
    end

    function foo.exit_default(self)
        print("leaving state", self.state)
    end
    
    function foo.exit_a(self)
        print("OMG I'm leaving A!")
    end
    
    foo.transition_a = {"a", "b", "c"}
    foo.transition_b = {"b", "a", "c"}
    foo.transition_c = {"c", "a", "b"}
    
    return foo()
end
